#include "Basis.h"
#include "KeyMouse.h"

KEYMOUSE_INFO g_KeyMouseInfo;

VOID InitializeKeyMouse()
{
	LookupServiceCallback();
}

VOID DisposeKeyMouse()
{
}

NTSTATUS LookupServiceCallback()
{
	NTSTATUS			status;
	UNICODE_STRING		ObjectName;

	PDRIVER_OBJECT		KeyboardClassDriverObject;
	PDEVICE_OBJECT		KeyboardClassDeviceObject;

	ULONG				KeyboardClassDriverStart;
	ULONG				KeyboardClassDriverEnd;

	PDRIVER_OBJECT		KeybaordPortDriverObject;
	PDEVICE_OBJECT		KeyboardPortDeviceObject;
	PDEVICE_OBJECT		KeyboardPortAttachedDeviceObject;
	PULONG				KeyboardPortDeviceExtension;
	ULONG				KeyboardPortDeviceExtensionSize;
	ULONG				i;

	// keyboard class driver
	RtlInitUnicodeString(&ObjectName, KEYBOARD_CLASS_NAME);
	status = ObReferenceObjectByName(&ObjectName, OBJ_CASE_INSENSITIVE, NULL, FILE_READ_ACCESS, *IoDriverObjectType, KernelMode, NULL, (PVOID*)&KeyboardClassDriverObject);
	if (!NT_SUCCESS(status))
	{
		KdPrint(("failed to reference class driver"));
		return status;
	}

	// class driver range
	KeyboardClassDriverStart = (ULONG)KeyboardClassDriverObject->DriverStart;
	KeyboardClassDriverEnd = (ULONG)KeyboardClassDriverStart + KeyboardClassDriverObject->DriverSize;
	KdPrint(("class driver start from 0x%x to 0x%x", KeyboardClassDriverStart, KeyboardClassDriverEnd));

	// keyboard port driver - usb
	RtlInitUnicodeString(&ObjectName, KEYBOARD_HID_NAME);
	status = ObReferenceObjectByName(&ObjectName, OBJ_CASE_INSENSITIVE, NULL, FILE_READ_ACCESS, *IoDriverObjectType, KernelMode, NULL, (PVOID*)&KeybaordPortDriverObject);
	if (!NT_SUCCESS(status))
	{
		// keyboard port driver - ps2
		RtlInitUnicodeString(&ObjectName, I8042PRT_NAME);
		status = ObReferenceObjectByName(&ObjectName, OBJ_CASE_INSENSITIVE, NULL, FILE_READ_ACCESS, *IoDriverObjectType, KernelMode, NULL, (PVOID*)&KeybaordPortDriverObject);
		if (!NT_SUCCESS(status))
		{
			KdPrint(("failed to reference port driver"));
			return status;
		}
	}

	// seeking callback
	KeyboardPortDeviceObject = KeybaordPortDriverObject->DeviceObject;
	while (KeyboardPortDeviceObject)
	{
		KeyboardPortAttachedDeviceObject = KeyboardPortDeviceObject;
		while (KeyboardPortAttachedDeviceObject)
		{
			KeyboardPortDeviceExtension = (PULONG)KeyboardPortAttachedDeviceObject->DeviceExtension;
			KeyboardPortDeviceExtensionSize = (ULONG)KeyboardPortAttachedDeviceObject->DeviceObjectExtension - (ULONG)KeyboardPortAttachedDeviceObject->DeviceExtension;
			KdPrint(("seeking at extension address 0x%x, size 0x%x(0x%x bytes)", KeyboardPortDeviceExtension, KeyboardPortDeviceExtensionSize / 4, KeyboardPortDeviceExtensionSize));

			KeyboardClassDeviceObject = KeyboardClassDriverObject->DeviceObject;
			while (KeyboardClassDeviceObject)
			{
				// do checking
				for (i = 0; i < KeyboardPortDeviceExtensionSize / 4; i++)
				{
					if ((ULONG)KeyboardPortDeviceExtension[i] == (ULONG)KeyboardClassDeviceObject)
					{
						KdPrint(("matched class device address 0x%x", KeyboardClassDeviceObject));

						if ((ULONG)KeyboardPortDeviceExtension[i + 1] > KeyboardClassDriverStart &&
							(ULONG)KeyboardPortDeviceExtension[i + 1] < KeyboardClassDriverEnd)
						{
							KdPrint(("matched class callback address 0x%x", KeyboardPortDeviceExtension[i + 1]));

							// save
							g_KeyMouseInfo.KeyboardDeviceObject = (PDEVICE_OBJECT)KeyboardPortDeviceExtension[i];
							g_KeyMouseInfo.KeyboardCallback = (KeyboardClassServiceCallback)KeyboardPortDeviceExtension[i + 1];

							KdPrint(("service callback found"));
							return STATUS_SUCCESS;
						}
					}
				}

				KeyboardClassDeviceObject = KeyboardClassDeviceObject->NextDevice; // next
			}

			KdPrint(("service callback not found"));
			KeyboardPortAttachedDeviceObject = KeyboardPortAttachedDeviceObject->AttachedDevice; // next
		}

		KeyboardPortDeviceObject = KeyboardPortDeviceObject->NextDevice; // next
	}

	return STATUS_UNSUCCESSFUL;
}

BOOLEAN ProcessIoCtlForKeyMouse(ULONG ioControlCode, PDEVICE_OBJECT deviceObject, PIRP irp, PIO_STACK_LOCATION irpStack)
{
	UNREFERENCED_PARAMETER(deviceObject);
	UNREFERENCED_PARAMETER(irpStack);

	if (ioControlCode != IOCTL_KEYBOARD) return FALSE;

	if (g_KeyMouseInfo.KeyboardCallback == NULL)
	{
		irp->IoStatus.Status = STATUS_UNSUCCESSFUL;
	}
	else
	{
		PKEYBOARD_INPUT_DATA pKeyboardInputData = (PKEYBOARD_INPUT_DATA)irp->AssociatedIrp.SystemBuffer;
		ULONG InputDataConsumed;
		g_KeyMouseInfo.KeyboardCallback(g_KeyMouseInfo.KeyboardDeviceObject, pKeyboardInputData, pKeyboardInputData + 1, &InputDataConsumed);
		KdPrint(("IOCTL_KEYBOARD consumed %d, flag 0x%x, code 0x%x", InputDataConsumed, pKeyboardInputData->Flags, pKeyboardInputData->MakeCode));

		irp->IoStatus.Status = STATUS_SUCCESS;
	}

	return TRUE;
}